from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, MessagesPlaceholder, \
    HumanMessagePromptTemplate
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory

from models import get_ds_model_client

client = get_ds_model_client()
prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template("你是一个聊天助手，用{language}回答所有的问题"),
    MessagesPlaceholder(variable_name="messages"),
    HumanMessagePromptTemplate.from_template("{input}")
])
parser = StrOutputParser()
chain = prompt | client | parser

# 会话存储
store = {}

# 不同session_id获得不同的会话历史记录ChatMessageHistory()
def get_session(session_id: str):
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


chatbot_with_history = RunnableWithMessageHistory(chain, get_session, input_messages_key="input",
                                                  history_messages_key="messages")

# configurable配置不同会话的session_id，这个id可以根据分布式ID生成数据
config_chn = {"configurable":{"session_id": "kenney_chn"}}
config_eng = {"configurable":{"session_id": "kenney_eng"}}

resp1 = chatbot_with_history.invoke({"input": "你好，我是kenney", "language": "中文"}, config=config_chn)
print(resp1)
resp2 = chatbot_with_history.invoke({"input": "你好，我是谁?", "language": "中文"}, config=config_chn)
print(resp2)

resp3 = chatbot_with_history.invoke({"input": "你好，我是kenney", "language": "英文"}, config=config_eng)
print(resp3)
resp4 = chatbot_with_history.invoke({"input": "你好，我是谁?", "language": "英文"}, config=config_eng)
print(resp4)